from typing import List

from bot.openai_bot import OpenAIBot
from db.mysql.strategy_title_dao import StrategyTitleDAO as StrategyTitleMysqlDAO
from db.milvus.strategy_title_dao import StrategyTitleDAO as StrategyTitleMIlvusDAO

class StrategyTitleService(object):
    """ 攻略标题服务对象 """

    def __init__(self, *args, **kwargs):
        self._strategy_title_milvus_dao = StrategyTitleMIlvusDAO()
        self._strategy_title_mysql_dao = StrategyTitleMysqlDAO()
        super(StrategyTitleService, self).__init__(*args, **kwargs)

    def get_title_ids(self, title_vector: List[float], top_k: int):
        """"""
        title_ids = self._strategy_title_milvus_dao.get_ids(title_vector=title_vector, top_k=top_k)
        return title_ids

    def get_title_tags(self, title_vector: List[float], top_k: int):
        title_ids = self.get_title_ids(title_vector=title_vector, top_k=top_k)
        for item in title_ids:
            title_id = item.get("title_id")
            title_info = self._strategy_title_mysql_dao.get_title_info(title_id=title_id)
            item.update({"title": title_info.get("title"), "title_tag": title_info.get("title_tag")})
        return title_ids


if __name__ == '__main__':
    strategy_title_service = StrategyTitleService()
    openai_bot = OpenAIBot()
    vector = openai_bot.embeddings(text_list=["大厂契税缴费标准", ])[0]
    # res = strategy_title_service.get_title_ids(title_vector=vector, top_k=10)
    res = strategy_title_service.get_title_tags(title_vector=vector, top_k=10)
    print(f"{res}")
